import sys
from numpy import ix_, fromfile, float64
import pylab
from pca import PCA
from dump import names, dim
from lsa import LSA

def add_names(nms, color='black'):
   nms_idx = [names.index(n) for n in nms]
   idx = ix_(nms_idx)
   for i in idx[0]:
      c = color
      pylab.text(d[i][0], d[i][1], names[i], horizontalalignment='center',
       verticalalignment='center', color=c)

def add_cluster(r, word, n=10, color='green'):
   names = [nm for type,idx,nm,score in r.find_closest(word, n) \
    if idx!=-1 and type==0]
   add_names(names, color)

# load data and do pca
with open("dump.dat", "rb") as f:
   a = fromfile(file=f, dtype=float64).reshape(-1, dim)
p = PCA(a)
p.npc = 2
d = p.pc()

if False:
   # plot points for all words
   pylab.plot(d[:,0], d[:,1], 'o')
else:
   # show some clusters, given a search name
   r = LSA()
   if len(sys.argv)>1:
      r.load(sys.argv[1])
      colors = ('green','red','blue','yellow','pink','black','gray')
      for i,nm in enumerate(sys.argv[3:10]):
         add_cluster(r,nm,10,colors[i])
   else:
      r.load('wiki.rav')
      add_cluster(r,'baseball',10,'green')
      add_cluster(r,'organism',10,'red')

pylab.gca().set_xlim(-1,1)
pylab.gca().set_ylim(-1,1)
if 'show' in sys.argv:
   pylab.show()
if 'save' in sys.argv:
   pylab.savefig('words.png')
